import numpy as np
import pandas as pd
import os


def get_randomdata_with_3sigma(mean, std, size):
    data = []
    while len(data) < size:
        samples = np.random.normal(loc=mean, scale=std, size=size)
        # 筛选满足 0 <= x <= mean + 3*std 的样本（同时排除负值）
        valid_samples = [x for x in samples if 0 <= x <= mean + 3 * std]
        data.extend(valid_samples)
    return np.array(data[:size])


def z_data(name, table):
    dataset = pd.read_excel(name, sheet_name=table)
    x = dataset.iloc[:, 0].values.astype(float)

    mean_x = np.mean(x)
    std_x = np.std(x)

    ex_size = 89

    # 使用 3σ 原则生成扩充数据并确保非负
    ex_x = get_randomdata_with_3sigma(mean_x, std_x, ex_size)
    all_x = ex_x  # 无需反标准化，已在原分布范围内生成

    df_data = pd.DataFrame({
        'x': x.tolist(),
    })

    df_all = pd.DataFrame({
        'x': all_x.tolist(),
    })

    return df_data, df_all


if __name__ == "__main__":
    input_path = os.path.join('数据', '数据.xlsx')

    if os.path.exists(input_path):
        df_data, df_all = z_data(input_path, 'Sheet1')
    else:
        print(f"错误：文件 {input_path} 不存在，请检查路径和文件名")

    os.makedirs('数据', exist_ok=True)
    output_path = os.path.join('数据', '数据扩充.xlsx')
    with pd.ExcelWriter(output_path, engine='xlsxwriter') as writer:
        df_data.to_excel(writer, sheet_name='标准化', index=False)
        df_all.to_excel(writer, sheet_name='整体数据', index=False)

    print("基于3σ原则的数据扩充已完成，负值已过滤。")
